package org.luxor.cloud.gateway.component.trace;

import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.cloud.gateway.support.DefaultServerRequest;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.util.Assert;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.List;

/**
 * 访问日志记录
 *
 * @author Mr.Yan  @date: 2020/11/8
 */
public class TraceGlobalFilter implements GlobalFilter, Ordered {
    private static final Logger logger = LoggerFactory.getLogger(TraceGlobalFilter.class);

    private TraceRegistry registry;

    private final List<MediaType> readable_media_types = Arrays.asList(MediaType.APPLICATION_JSON_UTF8
            , MediaType.APPLICATION_JSON
            , MediaType.APPLICATION_XML
            , MediaType.TEXT_XML
            , MediaType.TEXT_PLAIN);

    public TraceGlobalFilter(TraceRegistry registry) {
        Assert.notNull(registry, "Spring容器中找不到[RequestRecordService]的实例化对象!");
        this.registry = registry;
    }

    @Override
    public int getOrder() {
        return HIGHEST_PRECEDENCE + 10101;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        logger.trace("RequestRecordFilter start...");

        // 质量监控统计，计时开始
        Collector collector = Collector.start(registry);

        //如果支持读取请求/响应体内容(目前支持三种格式：application/json;application/xml;text/plain)
        if (readable_media_types.contains(exchange.getRequest().getHeaders().getContentType())) {
            return readBody(exchange, chain, collector).doOnSuccessOrError((aVoid, ex) -> {
                if (ex != null) {
                    logger.error("网关调用异常", ex);
                }
                endTimerRespectingCommit(exchange, collector);
                logger.trace("RequestRecordFilter end.");
            });
        } else {
            return chain.filter(exchange).doOnSuccessOrError((aVoid, ex) -> {
                if (ex != null) {
                    logger.error("网关调用异常", ex);
                }
                endTimerRespectingCommit(exchange, collector);
                logger.trace("RequestRecordFilter end.");
            });
        }

    }

    /**
     * 质量监控统计，计时结束
     */
    private void endTimerRespectingCommit(ServerWebExchange exchange, Collector collector) {
        ServerHttpResponse response = exchange.getResponse();
        if (response.isCommitted()) {
            collector.stop(exchange);
        } else {
            response.beforeCommit(() -> {
                collector.stop(exchange);
                return Mono.empty();
            });
        }
    }

    /**
     * 读取请求/响应体内容
     */
    private Mono<Void> readBody(ServerWebExchange exchange, GatewayFilterChain chain, Collector collector) {
        Mono<String> publisher = new DefaultServerRequest(exchange).bodyToMono(String.class)
                .flatMap(requestBody -> {
                    // 采集请求体内容
                    collector.setRequestBody(requestBody);
                    return Mono.just(requestBody);
                });

        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());

        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
        BodyInserter bodyInserter = BodyInserters.fromPublisher(publisher, String.class);
        return bodyInserter.insert(outputMessage, new BodyInserterContext())
                .then(Mono.defer(() -> {
                    // 请求包装类
                    ServerHttpRequestDecorator requestDecorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return outputMessage.getBody();
                        }
                    };

                    // 响应包装类
                    ServerHttpResponseDecorator responseDecorator = new ServerHttpResponseDecorator(exchange.getResponse()) {
                        @Override
                        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                            Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                            return super.writeWith(fluxBody.buffer().map(dataBuffers -> {

                                DataBuffer join = new DefaultDataBufferFactory().join(dataBuffers);
                                byte[] content = new byte[join.readableByteCount()];
                                join.read(content);
                                DataBufferUtils.release(join);
                                String responseBody = new String(content, Charset.forName("UTF-8"));
                                // 采集响应体内容
                                collector.setResponseBody(responseBody);
                                exchange.getResponse().getHeaders().setContentLength(responseBody.getBytes().length);
                                return exchange.getResponse().bufferFactory().wrap(responseBody.getBytes());
                            }));
                        }
                    };

                    // 继续执行
                    return chain.filter(exchange.mutate().request(requestDecorator).response(responseDecorator).build());
                }));
    }
}